from absl import app, flags

import pandas
import datetime
import functools
import logging
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from coin.base.book.types import TradeSide
from coin.exchange.okex.kr_rest.product import OkexProduct
from coin.exchange.huobi.kr_rest.product import HuobiProduct
from coin.exchange.shared.kr_rest.product import SharedProduct
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.bitflyer_v1.kr_rest.futures_product import BitflyerFuturesProduct
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str

from coin.strategy.mm.dumper_base import run_from_archive
from coin.support.proto_log.logic.util import run_from_og_log_archive
import coin.strategy.mm.feed as feed
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog, OrderEvent
from coin.strategy.mm.tool.util import get_order_sign, KuonaOrderManager
from coin.strategy.tool.monitor.dump_plot import query_balance_from_log
from coin.strategy.tool.monitor.dump_plot import _aggregate_by_key

FLAGS = flags.FLAGS

import coin.strategy.mm.tool.plot_util as putil


class Plotter(object):
  def __init__(self, feed_product, true_book_qty, filename):
    self.feed_product = feed_product
    self.obs = []
    self.order_book = {}
    self.true_book_qty = true_book_qty
    self.filename = filename

  def on_book_reset(self, book_builder_name, book_builder):
    for product in [self.feed_product]:
      subscription_key = feed.get_subscription_key(product)
      if book_builder_name == subscription_key:
        book_builder.subscribe(product, functools.partial(self.on_book, product))
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def on_trade(self, product, trade):
    if product != self.feed_product:
      return
    self.order_book['time'] = trade.timestamp
    if trade.side == -1:
      self.order_book['stradep'] = trade.price
    else:
      self.order_book['btradep'] = trade.price
    self.order_book['tradeq'] = trade.qty
    self.obs.append(self.order_book.copy())
    self.order_book['stradep'] = None
    self.order_book['btradep'] = None
    self.order_book['tradeq'] = None

  def on_book(self, product, book):
    if product != self.feed_product:
      return
    self.order_book['time'] = book.timestamp
    self.order_book['bid0p'] = book.get_bid_price_by_qty(self.true_book_qty)
    self.order_book['ask0p'] = book.get_ask_price_by_qty(self.true_book_qty)
    self.obs.append(self.order_book.copy())

  def plot(self, result_given):
    df = pandas.DataFrame(self.obs)
    df.index = pandas.DatetimeIndex(df['time'])
    matplotlib.style.use('fivethirtyeight')
    matplotlib.rcParams.update({'font.size': 7})
    plt.subplot(211)
    plt.title("%s, true_book_qty: %s" % (self.feed_product, self.true_book_qty))
    plt.plot(df['bid0p'], 'g-', df['ask0p'], 'g-', linewidth=0.1)
    plt.plot(df['btradep'], 'g.', df['stradep'], 'r.', markersize=0.1)
    plt.ylabel('price')
    if result_given is not None:
      plt.twinx()
      plt.grid(None)
      for result in result_given:
        baldf = pandas.DataFrame(result['balance_dump'])
        baldf.index = pandas.DatetimeIndex(baldf['time'])
        plt.plot(baldf['value'] - baldf['value'].iloc[0], 'k-', linewidth=0.5)
      plt.ylabel('pnl')
    plt.subplot(413)
    # spdplt = plt.plot(df['ask0p'] - df['bid0p'], 'g-', linewidth=0.2, label='spd')
    # rollspd = pandas.Series(df['ask0p'] - df['bid0p']).rolling(300).mean()
    spd_bps = 1e4 * (df['ask0p'] - df['bid0p']) / df['ask0p']
    twindow = '5min'
    twindow2 = '30min'
    rollspd = pandas.Series(spd_bps).rolling(twindow, min_periods=30).mean()
    rollspdplt = plt.plot(rollspd, 'r-', linewidth=0.2, label='rollspd_bps_%s' % twindow)
    rollspd2 = pandas.Series(spd_bps).rolling(twindow2, min_periods=30).mean()
    rollspdplt2 = plt.plot(rollspd2, 'm-', linewidth=0.2, label='rollspd_bps_%s' % twindow2)
    plt.ylabel('spread(bps)')
    plt.legend(handles=rollspdplt + rollspdplt2)
    plt.subplot(414)
    rolltradeq = pandas.Series(df['tradeq']).rolling(twindow, min_periods=30).sum()
    rolltrqplt = plt.plot(rolltradeq, 'g-', linewidth=0.2, label='roll_tradeq_%s' % twindow)
    rolltradeq2 = pandas.Series(df['tradeq']).rolling(twindow2, min_periods=30).sum()
    rolltrqplt2 = plt.plot(rolltradeq2, 'c-', linewidth=0.2, label='roll_tradeq_%s' % twindow2)
    plt.ylabel('trade qty')
    plt.legend(handles=rolltrqplt + rolltrqplt2)
    plt.tight_layout()
    plt.savefig(self.filename, dpi=300)
    plt.close()


def launch(trading_date, exchange, symbol):
  start_time = trading_date  # UTC 0 hour
  actual_exchange = exchange
  if FLAGS.market_type == 'Futures' and FLAGS.exchange == 'Okex':
    product = OkexFuturesProduct.FromStr(symbol, start_time)
  else:
    product = generate_product_from_str(actual_exchange, None, symbol)
  add_hour = 0
  hours = 12
  # hours = 2
  i = 0
  while add_hour < 24:
    i = i + 1

    if FLAGS.pnl_group is not None:
      results = query_balance_from_log(start_time, start_time + datetime.timedelta(hours=hours))
      results_by_group = _aggregate_by_key(results, 'group')
      result_chosen = results_by_group[FLAGS.pnl_group]
    else:
      result_chosen = None

    strategy = Plotter(product,
                       FLAGS.true_book_qty,
                       "pic/%s_%s_%s.png" % (symbol, FLAGS.trading_date, i))
    run_from_archive([strategy.feed_product],
                     strategy.on_book_reset,
                     start_time,
                     start_time + datetime.timedelta(hours=hours),
                     machine=[FLAGS.feed_machine])
    strategy.plot(result_chosen)
    start_time += datetime.timedelta(hours=hours)
    add_hour += hours
  return 0


def main(_):
  trading_date = FLAGS.trading_date
  exchange = FLAGS.exchange
  symbol = FLAGS.symbol
  assert FLAGS.market_type in ['Spot', 'Futures']
  assert trading_date, '--trading_date must be specified.'
  assert exchange, '--exchange must be specified.'
  print('Running for %s %s %s ...' % (trading_date, symbol, exchange))
  sys.stdout.flush()
  trading_date = datetime.datetime.strptime(trading_date, '%Y%m%d')
  launch(trading_date, exchange, symbol)


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('pnl_group', None, '')

  flags.DEFINE_string('feed_machine', 'feed-01.eu-west-1.aws', 'feed machine name.')

  flags.DEFINE_string('exchange', 'Okex', 'Exchange name.')

  flags.DEFINE_string('market_type', 'Spot', 'Spot, Futures')

  flags.DEFINE_string('symbol', 'BTC-USDT', 'symbol name')

  flags.DEFINE_float('true_book_qty', 0, 'true book qty')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
